#!/usr/bin/env python3
import argparse
import os
import sys
from time import sleep

import grpc

# Import P4Runtime lib from parent utils dir
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections

def writeIpv4Rules(p4info_helper, switch, switch_name):
    """
    Installs IPv4 forwarding rules for the switch based on topology

    :param p4info_helper: the P4Info helper
    :param switch: the switch connection
    :param switch_name: name of the switch (s1, s2, s3)
    """

    # Define routing tables for each switch based on topology
    if switch_name == 's1':
        # s1 routing rules
        rules = [
            # dst_ip/mask, dst_mac, port
            ("10.0.1.0", 24, "08:00:00:00:01:11", 1),  # h1
            ("10.0.2.0", 24, "08:00:00:00:02:00", 2),  # s2
            ("10.0.3.0", 24, "08:00:00:00:03:00", 3),  # s3
            ("10.0.1.1", 32, "08:00:00:00:01:11", 1),  # h1 specific
            ("10.0.2.2", 32, "08:00:00:00:02:00", 2),  # h2 via s2
            ("10.0.3.3", 32, "08:00:00:00:03:00", 3),  # h3 via s3
        ]
    elif switch_name == 's2':
        # s2 routing rules
        rules = [
            ("10.0.1.0", 24, "08:00:00:00:01:00", 2),  # s1
            ("10.0.2.0", 24, "08:00:00:00:02:22", 1),  # h2
            ("10.0.3.0", 24, "08:00:00:00:03:00", 3),  # s3
            ("10.0.1.1", 32, "08:00:00:00:01:00", 2),  # h1 via s1
            ("10.0.2.2", 32, "08:00:00:00:02:22", 1),  # h2 specific
            ("10.0.3.3", 32, "08:00:00:00:03:00", 3),  # h3 via s3
        ]
    elif switch_name == 's3':
        # s3 routing rules
        rules = [
            ("10.0.1.0", 24, "08:00:00:00:01:00", 2),  # s1
            ("10.0.2.0", 24, "08:00:00:00:02:00", 3),  # s2
            ("10.0.3.0", 24, "08:00:00:00:03:33", 1),  # h3
            ("10.0.1.1", 32, "08:00:00:00:01:00", 2),  # h1 via s1
            ("10.0.2.2", 32, "08:00:00:00:02:00", 3),  # h2 via s2
            ("10.0.3.3", 32, "08:00:00:00:03:33", 1),  # h3 specific
        ]
    else:
        print(f"Unknown switch: {switch_name}")
        return

    print(f"\n----- Installing IPv4 forwarding rules for {switch_name} -----")

    for dst_ip, prefix_len, dst_mac, port in rules:
        table_entry = p4info_helper.buildTableEntry(
            table_name="MyIngress.ipv4_lpm",
            match_fields={
                "hdr.ipv4.dstAddr": (dst_ip, prefix_len)
            },
            action_name="MyIngress.ipv4_forward",
            action_params={
                "dstAddr": dst_mac,
                "port": port
            })
        try:
            switch.WriteTableEntry(table_entry)
            print(f"Installed rule: {dst_ip}/{prefix_len} -> MAC: {dst_mac}, Port: {port}")
        except Exception as e:
            print(f"Failed to install rule for {dst_ip}/{prefix_len}: {e}")

def writeMacLearningRules(p4info_helper, switch, switch_name):
    """
    Installs MAC learning rules for ARP resolution
    """
    if switch_name == 's1':
        mac_entries = [
            ("08:00:00:00:01:11", 1),  # h1
            ("08:00:00:00:02:00", 2),  # s2 direction
            ("08:00:00:00:03:00", 3),  # s3 direction
        ]
    elif switch_name == 's2':
        mac_entries = [
            ("08:00:00:00:02:22", 1),  # h2
            ("08:00:00:00:01:00", 2),  # s1 direction
            ("08:00:00:00:03:00", 3),  # s3 direction
        ]
    elif switch_name == 's3':
        mac_entries = [
            ("08:00:00:00:03:33", 1),  # h3
            ("08:00:00:00:01:00", 2),  # s1 direction
            ("08:00:00:00:02:00", 3),  # s2 direction
        ]
    else:
        return

    print(f"\n----- Installing MAC entries for {switch_name} -----")

    # Note: This is optional since we're using IPv4 forwarding with MAC rewriting
    # But it's good practice to have MAC learning capabilities
    for mac_addr, port in mac_entries:
        print(f"MAC: {mac_addr} -> Port: {port}")

def readTableRules(p4info_helper, switch):
    """
    Reads the table entries from all tables on the switch.

    :param p4info_helper: the P4Info helper
    :param switch: the switch connection
    """
    print(f'\n----- Reading table rules for {switch.name} -----')

    try:
        # Read IPv4 forwarding table
        for response in switch.ReadTableEntries():
            for entity in response.entities:
                entry = entity.table_entry
                table_name = p4info_helper.get_name('tables', entry.table_id)

                if table_name == "MyIngress.ipv4_lpm":
                    # Extract match fields
                    match_fields = {}
                    for mf in entry.match:
                        field_name = p4info_helper.get_name('match_fields', mf.field_id)
                        if mf.lpm:
                            match_fields[field_name] = f"{mf.lpm.value}/{mf.lpm.prefix_len}"
                        elif mf.exact:
                            match_fields[field_name] = mf.exact.value

                    # Extract action
                    action = entry.action.action
                    action_name = p4info_helper.get_name('actions', action.action_id)
                    action_params = {}

                    for param in action.params:
                        param_name = p4info_helper.get_name('action_params', param.param_id)
                        action_params[param_name] = param.value

                    print(f"Table: {table_name}")
                    print(f"  Match: {match_fields}")
                    print(f"  Action: {action_name} {action_params}")
                    print('-----')

    except Exception as e:
        print(f"Error reading table rules: {e}")

def printCounter(p4info_helper, switch, counter_name, index=0):
    """
    Reads the specified counter from the switch.

    :param p4info_helper: the P4Info helper
    :param switch: the switch connection
    :param counter_name: the name of the counter from the P4 program
    :param index: the counter index
    """
    try:
        counter_id = p4info_helper.get_counters_id(counter_name)
        for response in switch.ReadCounters(counter_id, index):
            for entity in response.entities:
                counter = entity.counter_entry
                print(f"{switch.name} {counter_name}[{index}]: {counter.data.packet_count} packets, {counter.data.byte_count} bytes")
    except Exception as e:
        print(f"Error reading counter {counter_name}: {e}")

def main(p4info_file_path, bmv2_file_path):
    # Instantiate a P4Runtime helper from the p4info file
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # Create switch connection objects for s1, s2, s3
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2',
            address='127.0.0.1:50052',
            device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3',
            address='127.0.0.1:50053',
            device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        # Send master arbitration update message to establish this controller as master
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        # Install the P4 program on the switches
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s1")
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s2")
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s3")

        # Write IPv4 forwarding rules for each switch
        writeIpv4Rules(p4info_helper, s1, 's1')
        writeIpv4Rules(p4info_helper, s2, 's2')
        writeIpv4Rules(p4info_helper, s3, 's3')

        # Write MAC learning rules (optional)
        writeMacLearningRules(p4info_helper, s1, 's1')
        writeMacLearningRules(p4info_helper, s2, 's2')
        writeMacLearningRules(p4info_helper, s3, 's3')

        # Read and display installed table rules
        print("\n=== Verifying installed rules ===")
        readTableRules(p4info_helper, s1)
        readTableRules(p4info_helper, s2)
        readTableRules(p4info_helper, s3)

        # Monitor counters periodically
        print("\n=== Starting counter monitoring (Ctrl+C to stop) ===")
        while True:
            sleep(5)
            print('\n----- Reading packet counters -----')
            # Note: You can add counters to your P4 program if needed
            # printCounter(p4info_helper, s1, "MyIngress.packet_counter")
            # printCounter(p4info_helper, s2, "MyIngress.packet_counter")
            # printCounter(p4info_helper, s3, "MyIngress.packet_counter")

    except KeyboardInterrupt:
        print("\nShutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller for Source Routing with IPv4')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/source_routing.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/source_routing.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)
    main(args.p4info, args.bmv2_json)
